// 我们有一个由平面上的点组成的列表 points。需要从中找出 K 个距离原点 (0, 0) 最近的点。

// （这里，平面上两点之间的距离是欧几里德距离。）

// 你可以按任何顺序返回答案。除了点坐标的顺序之外，答案确保是唯一的。

//  

// 示例 1：

// 输入：points = [[1,3],[-2,2]], K = 1
// 输出：[[-2,2]]
// 解释： 
// (1, 3) 和原点之间的距离为 sqrt(10)，
// (-2, 2) 和原点之间的距离为 sqrt(8)，
// 由于 sqrt(8) < sqrt(10)，(-2, 2) 离原点更近。
// 我们只需要距离原点最近的 K = 1 个点，所以答案就是 [[-2,2]]。
// 示例 2：

// 输入：points = [[3,3],[5,-1],[-2,4]], K = 2
// 输出：[[3,3],[-2,4]]
// （答案 [[-2,4],[3,3]] 也会被接受。）
//  

// 提示：

// 1 <= K <= points.length <= 10000
// -10000 < points[i][0] < 10000
// -10000 < points[i][1] < 10000

#include "stdc++.h"

/* 排序
时间复杂度：O(NlogN)
空间复杂度：O(N)
*/
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        int n = points.size();
        vector<int> dists(n, 0);
        for (int i{0}; i < n; ++i) {
            dists[i] = points[i][0] * points[i][0] + points[i][1] * points[i][1];
        }
        sort(dists.begin(), dists.end());
        int distK = dists[K - 1]; // 第K小的距离
        vector<vector<int>> res(K, vector<int>(2, 0));
        int i{0};
        for (const auto& point : points) {
            if (point[0] * point[0] + point[1] * point[1] <= distK) { // 比第K小的距离小的都是我们要的，不需要考虑顺序
                res[i++] = point;
            }
        }
        return res;
    }
};

/* 分治
时间复杂度：O(N)
空间复杂度：O(N)
*/
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        this->points = points;
        work(0, this->points.size() - 1, K);
        return vector<vector<int>>(this->points.begin(), this->points.begin() + K);
    }
    // 分治
    void work(int i, int j, int K) {
        if (i >= j) {
            return;
        }
        int oi{i};
        int oj{j};
        int pivot = dist(i);
        while (i < j) {
            while (i < j && dist(i) < pivot) ++i;
            while (i < j && dist(j) > pivot) --j;
            if (dist(i) == pivot && dist(j) == pivot) break;
            swap(i, j);
        }
        if (K <= i - oi + 1) {
            work(oi, i, K);
        } else {
            work(i + 1, oj, K - (i - oi + 1));
        }
    }
    int dist(int i) {
        return points[i][0] * points[i][0] + points[i][1] * points[i][1];
    }
    void swap(int i, int j) {
        int t0 = points[i][0];
        int t1 = points[i][1];
        points[i][0] = points[j][0];
        points[i][1] = points[j][1];
        points[j][0] = t0;
        points[j][1] = t1;
    }
private:
    vector<vector<int>> points{};
};

#include "stdc++.h"

/* 排序
*/
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        auto cmp = [] (const vector<int>& pA, const vector<int>& pB) {
            return pA[0] * pA[0] + pA[1] * pA[1] < pB[0] * pB[0] + pB[1] * pB[1]; 
        };
        sort(points.begin(), points.end(), cmp);
        return {points.begin(), points.begin() + K};
    }
};

/* 优先队列
pair里第二个元素是 数组的下标
*/
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        priority_queue<pair<int, int>> pq{};
        for (int i{0}; i < K; ++i) {
            pq.emplace(points[i][0] * points[i][0] + points[i][1] * points[i][1], i);
        }
        int n = points.size();
        for (int i{K}; i < n; ++i) {
            int dist = points[i][0] * points[i][0] + points[i][1] * points[i][1];
            if (dist < pq.top().first) {
                pq.pop();
                pq.emplace(dist, i);
            }
        }
        vector<vector<int>> res{};
        while (!pq.empty()) {
            res.push_back(points[pq.top().second]);
            pq.pop();
        }
        return res;
    }
};

/* 快速排序
尾递归
*/
class Solution {
private:
    // 产生随机数
    mt19937 gen{
        random_device{}()
    };
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        int n = points.size();
        random_select(points, 0, n - 1, K);
        return {points.begin(), points.begin() + K};
    }

    void random_select(vector<vector<int>>& points, int left, int right, int K) { 
        int pivot_id = uniform_int_distribution<int>{left, right}(gen);
        int pivot = points[pivot_id][0] * points[pivot_id][0] +
                    points[pivot_id][1] * points[pivot_id][1];
        swap(points[right], points[pivot_id]);
        int i = left - 1;
        for (int j{left}; j < right; ++j) {
            int dist = points[j][0] * points[j][0] + points[j][1] * points[j][1];
            if (dist <= pivot) {
                ++i;
                swap(points[i], points[j]);
            }
        }
        ++i;
        swap(points[i], points[right]);
        // [left, i - 1] 都小于等于 pivot, [i + 1, right] 都大于 pivot
        if (K < i - left + 1) {
            random_select(points, left, i - 1, K);
        } else if (K > i - left + 1) {
            random_select(points, i + 1, right, K - (i - left + 1));
        }
    }
};